"""
Visdom可视化工具封装类

核心功能：
1. 实时监控训练过程中的指标变化（损失、准确率等）
2. 可视化模型输入/输出图像
3. 保存训练日志和可视化结果
4. 支持断点续训的状态保存与加载

技术特性：
- 基于Facebook开源的Visdom库构建
- 支持文本、图像、曲线等多种数据类型
- 自动维护可视化窗口的状态
- 可扩展的保存系统（日志/图像/曲线）
"""
from datetime import datetime

import visdom
import numpy as np
import time
import torch
import os

from PIL import Image

from config.config import Config as cfg



class Visualizer(object):
    """Visdom可视化工具封装类

    初始化参数：
        env (str): Visdom环境名称（默认'default'）
        **kwargs: Visdom服务器连接参数

    功能模块：
        - 文本日志记录
        - 图像可视化
        - 曲线绘制
        - 状态持久化
    """

    def __init__(self, env='default', **kwargs):
        """初始化Visdom连接和状态记录"""
        self.save_plot = None
        self.save_img = None
        self.save_log = None
        self.save_path = None
        self.env = cfg.vis_env # 可视化环境名称
        self.vis = visdom.Visdom(env=env, use_incoming_socket=True, **kwargs)  # Visdom客户端实例
        if not self.vis.check_connection():
            raise RuntimeError("Visdom服务器未启动！请先运行: python -m visdom.server")
        else:
            print("Visdom连接成功！等待训练开始...")
        self._vis_kwargs = kwargs  # 连接参数备份

        # 状态记录字典
        self.index = {}  # 记录曲线图的X轴坐标
        self.log_text = {
            'default': '421',
        }  # 文本日志缓存+


    def save_settings(self, save_path=None, save_log=False, save_img=False, save_plot=False):
        """配置保存选项

        参数：
            save_path (str): 保存根目录
            save_log (bool): 是否保存文本日志
            save_img (bool): 是否保存可视化图像
            save_plot (bool): 是否保存曲线数据

        目录结构：
            save_path/
                YYYY-MM-DD-HH-MM-SS_env/
                    logs/       # 文本日志
                    imgs/       # 图像快照
                    plots/      # 曲线数据
        """
        save_format = '{info}_{time}'.format(time=time.strftime("%Y_%m_%d_%H_%M_%S"), info=self.env)

        # print(save_format)
        self.save_path = os.path.join(save_path, save_format)
        # print("保存路径：")
        # print(self.save_path)

        # 创建保存目录
        if self.save_path and not os.path.exists(self.save_path):
            os.makedirs(self.save_path,exist_ok=True)

        # 初始化各子目录
        self.save_log = save_log
        if self.save_path and self.save_log:
            os.mkdir((os.path.join(self.save_path, 'logs')))

        self.save_img = save_img
        if self.save_path and self.save_img:
            os.mkdir((os.path.join(self.save_path, 'imgs')))

        self.save_plot = save_plot
        if self.save_path and self.save_plot:
            os.mkdir((os.path.join(self.save_path, 'plots')))

    def __getattr__(self, name):
        """透明化访问Visdom原生方法"""
        return getattr(self.vis, name)

    def reinit(self, env, **kwargs):
        """重新初始化Visdom连接"""
        self.vis = visdom.Visdom(env=env, **kwargs)

    def state_dict(self):
        """获取当前状态（用于断点续训）"""
        return {
            'index': self.index,
            'log_text': self.log_text,
            '_vis_kwargs': self._vis_kwargs,
            'env': self.vis.env,
        }

    def load_state_dict(self, d):
        """加载保存的状态"""
        self.vis = visdom.Visdom(env=d.get('env', self.vis.env),
                                 use_incoming_socket=False,
                                 **d.get('_vis_kwargs'))
        self.log_text = d.get('log_text', {})
        self.index = d.get('index', {})

    # 记录文本日志（同时保存到Visdom和本地文件）  edit by likang
    def log(self, info, win='default'):
        """
        记录文本日志（同时保存到Visdom和本地文件）
        """
        # 可视化到Visdom
        if self.log_text.get(win, None) is not None:
            flag = True  # 追加模式
        else:
            flag = False  # 新建窗口
        log_str = '[{time}] {info}\n'.format(time=time.strftime("%Y-%m-%d %H:%M:%S"), info=info)
        self.log_text[win] = log_str  # 缓存日志
        self.vis.text(log_str, win, append=flag)  # 发送到Visdom

        # print(self.save_log)
        # 保存到本地文件（如果启用了save_log）
        if self.save_log and self.save_path:
            log_file = os.path.join(self.save_path, 'logs', f'{win}.log')
            with open(log_file, 'a' if flag else 'w',  encoding='utf-8') as f:
                f.write(log_str)

    # 多次调用log  即:上面的log
    def log_many(self, d):
        """批量记录多组日志"""
        for k, v in d.items():
            self.log(v, k)

    #
    def img(self, img, win='default',idx=None, **kwargs):
        """可视化图像数据

        支持输入：
            - PyTorch Tensor (C,H,W)或(H,W)
            - NumPy数组
            - 批量图像 (B,C,H,W)
        """
        if torch.is_tensor(img):
            img = img.detach().cpu().numpy()  # 转换到CPU NumPy
        self.vis.images(img, win=win, opts=dict(title=win), **kwargs)
        save_path = os.path.join(self.save_path, 'imgs')
        date = datetime.now().strftime('%Y%m%d')
        if self.save_img and self.save_path:
            if img.ndim == 3:  # 单张图像 (C,H,W)
                img = img.transpose(1, 2, 0)  # 转为 (H,W,C)
                if img.dtype == np.float32:  # 处理float32类型
                    img = (img * 255).clip(0, 255).astype(np.uint8)  # 映射到[0,255]并转为uint8
                Image.fromarray(img).save(save_path+f"_{win+'-'+str(idx)}.png")
            elif img.ndim == 4:  # 批量图像 (B,C,H,W)
                for i, im in enumerate(img):
                    im = im.transpose(1, 2, 0)  # (C,H,W) -> (H,W,C)
                    if im.dtype == np.float32:  # 处理float32类型
                        im = (im * 255).clip(0, 255).astype(np.uint8)
                    Image.fromarray(im).save(f"{save_path}_{win+'-'+str(idx)}.png")

    def img_many(self, d,idx):
        # print(d)
        """批量可视化多组图像"""
        for k, v in d.items():
            # print('键',k)
            self.img(v, k,idx=idx)

    # 曲线日志  并保存csv数据到本地  csv手动转图片可查看
    def plot(self, y, win='loss', idx=None, **kwargs):
        """绘制单值曲线（同时保存坐标数据到本地）"""
        x = self.index.get(win, 0)  # 获取当前X坐标

        # 可视化到Visdom
        self.vis.line(
            Y=np.array([y]),
            X=np.array([x]),
            win=win,
            opts=dict(title=win),
            update=None if x == 0 else 'append',  # 首次创建/后续追加
            **kwargs
        )
        self.index[win] = x + 1  # 更新X轴索引

        # 保存到本地文件（如果启用了save_plot）
        if self.save_plot and self.save_path:
            # 生成文件路径：save_path/plots/{win}.csv
            plot_file = os.path.join(self.save_path, 'plots', f"{win+'-'+str(idx)}.csv")

            # 确定文件打开模式：首次写入（x=0）时创建文件并添加header，后续追加
            mode = 'w' if x == 0 else 'a'

            # 写入数据（x和y格式化为字符串，兼容CSV）
            with open(plot_file, mode, encoding='utf-8') as f:
                if mode == 'w':  # 首次写入时添加表头
                    f.write("step,value\n")
                f.write(f"{x},{y:.6f}\n")  # 保留6位小数，便于后续分析

    # 多次调用plot
    def plot_many(self, d,  idx):
        """批量绘制多条曲线"""
        for k, v in d.items():
            if v is not None:
                self.plot(v, k, idx=idx)



# 测试连接
if __name__ == "__main__":
    # 先检测连接
    con = Visualizer().vis.check_connection()
    print("Visdom连接成功！开始可视化演示...")
    visdom = Visualizer()
    visdom.save_settings(save_path='../log/visdom', save_log=True, save_img=True, save_plot=True)
    visdom.log("测试文本日志！")